# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: inkyu

import datetime
import hashlib
import logging
import random
import sys

import kafka
from absl import app, flags
from tornado.ioloop import IOLoop

from coin.experimental.inkyu.parsers.kimchi import KimchiParser
from coin.experimental.inkyu.parsers.unhedge import UnhedgeParser
from coin.experimental.inkyu.parsers.json import JsonParser
from coin.experimental.inkyu.parsers.hk import HkParser


def fg(color):
  return '\033[38;5;%dm' % color


strategies = {
    'strategy-01.ap-northeast-2_log.passagg_upbit_okex': (KimchiParser, 'upbit_okex'),
    'strategy-02.ap-northeast-2_log.unhedge_upbit_a': (UnhedgeParser, 'upbit_a'),
    'strategy-02.ap-northeast-2_log.unhedge_upbit_b': (UnhedgeParser, 'upbit_b'),
    'strategy-04.ap-northeast-2_log.unhedge_bithumb_b': (UnhedgeParser, 'bithumb_b'),
    'strategy-05.ap-northeast-2_log.passagg_bithumb_okex': (KimchiParser, 'bithumb_okex'),
    'strategy-06.ap-northeast-2_log.passagg_bitfinex_bithumb': (KimchiParser, 'bitfinex_bithumb'),
    'strategy-07.ap-northeast-2_log.passagg_huobi_bithumb': (KimchiParser, 'huobi_bithumb'),
    'strategy-08.ap-northeast-2_log.unhedge_bithumb_a': (UnhedgeParser, 'bithumb_a'),
    'strategy-10.ap-northeast-2_log.passagg_bitfinex_upbit': (KimchiParser, 'bitfinex_upbit'),
    'strategy-05.us-east-2_log.passagg_kraken_okex_eur': (JsonParser, 'kraken_okex_eur'),
    'strategy-03.ap-northeast-1.huobi_log.passagg_huobi_okex': (HkParser, 'huobi_okex'),
    'strategy-03.ap-northeast-1.huobi_log.passagg_huobi_binance': (HkParser, 'huobi_binance'),
    'strategy-04.ap-northeast-1.huobi_log.passagg_huobi_upbit': (KimchiParser, 'huobi_upbit'),
    'strategy-17.aliyun-cn-hongkong_log.passagg_hitbtc_okex': (HkParser, 'hitbtc_okex'),
    'strategy-05.ap-northeast-1.huobi_log.passagg_hitbtc_binance': (HkParser, 'hitbtc_binance'),
}


class KafkaMessagePusher(object):
  def __init__(self,
               *,
               kafka_servers,
               strategies=None,
               dry_run=True,
               color=True,
               randomize_color=False,
               logger=None):
    self._logger = logger or logging.getLogger(__name__)
    self._strategies = strategies
    self._kafka_servers = kafka_servers
    self._kafka_consumer = None
    self._ioloop = None
    self._dry_run = dry_run

    self._color = color
    self._color_offset = random.randint(0, 255) if randomize_color else 0

    self._observed_max_topic_len = 0
    self._parsers = {}

  async def _start(self):
    self._kafka_consumer = kafka.KafkaConsumer(bootstrap_servers=self._kafka_servers,
                                               auto_offset_reset='latest')
    for topic in self._kafka_consumer.topics():
      print(topic)
    assert self._strategies is not None

    self._logger.info('Topics: %s', ', '.join(self._strategies.keys()))
    self._kafka_consumer.subscribe(self._strategies.keys())

  async def _loop(self):
    while True:
      msg = next(self._kafka_consumer)
      topic = msg.topic
      if topic not in self._parsers:
        parser_cls, db_name = strategies[topic]
        self._parsers[topic] = parser_cls(db_name, dry_run=self._dry_run, on_push=self._on_push)
      parser = self._parsers[topic]
      for line in msg.value.decode('utf-8').split('\n'):
        parser.handle(line)

  def start(self, *, ioloop):
    self._ioloop = ioloop
    self._ioloop.run_sync(self._start)
    self._ioloop.add_callback(self._loop)

  def _topic_to_color_str(self, topic):
    hashed_topic = hashlib.md5(bytes(topic, 'ascii')).digest()
    color_no = (hashed_topic[4] + self._color_offset) % 256
    if color_no in (0, 15, 16, 231, 232, 233, 234, 235, 254, 255):
      color_no = (color_no + 87) % 256
    return fg(color_no)

  def _on_push(self, db_name, target, msg):
    dt_str = datetime.datetime.now().strftime('%m%d %H:%M:%S.%f')[:-3]
    topic_len = 16
    if self._color:
      color_str = self._topic_to_color_str(db_name)
      format_str = '%s%%-%ds%s │ %s%%17s%s │ %s%%17s%s ┃ %s%%s%s' % (color_str,
                                                                     topic_len,
                                                                     '\033[0m',
                                                                     color_str,
                                                                     '\033[0m',
                                                                     color_str,
                                                                     '\033[0m',
                                                                     color_str,
                                                                     '\033[0m')
    else:
      format_str = '%%-%ds │ %%17s │ %%17s ┃ %%s' % topic_len
    print(format_str % (db_name, dt_str, target, msg))


def main(argv):
  FLAGS = flags.FLAGS

  ioloop = IOLoop()
  msg_printer = KafkaMessagePusher(kafka_servers=FLAGS.kafka_servers.split(','),
                                   color=(sys.stdout.isatty()),
                                   strategies=strategies,
                                   dry_run=FLAGS.dry_run,
                                   randomize_color=False)
  msg_printer.start(ioloop=ioloop)

  try:
    ioloop.start()
  except (KeyboardInterrupt, SystemExit):
    print()
    return


if __name__ == '__main__':
  logging.basicConfig(
      level='DEBUG',
      format='%(levelname)8s %(asctime)s %(name)s %(filename)s:%(lineno)d] %(message)s')

  flags.DEFINE_string('kafka_servers',
                      'coin-kafka.corp.prestolabs.io:9092',
                      'Kafka server list separated by comma.')
  flags.DEFINE_boolean('dry_run', True, 'Dry run')

  app.run(main)
